import datetime
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.qastockbase import QAStrategyStockBase


class TurtleTrade(QAStrategyStockBase):
    '''
    本策略通过计算CZCE.FG801和SHFE.rb1801的ATR.唐奇安通道和MA线,并:
    上穿唐奇安通道且短MA在长MA上方则开多仓,下穿唐奇安通道且短MA在长MA下方则开空仓
    若有 多/空 仓位则分别:
    价格 跌/涨 破唐奇安平仓通道 上/下 轨则全平仓位,否则
    根据 跌/涨 破持仓均价 -/+ x(x=0.5,1,1.5,2)倍ATR把仓位
    回测数据为:CZCE.FG801和SHFE.rb1801的1min数据
    回测时间为:2017-09-15 09:15:00到2017-10-01 15:00:00
    '''

    def __init__(self, username, password, code='510300', start='2019-01-01', end='2019-10-21', show_trade=True,
                 frequence='1min', strategy_id='QA_STRATEGY', init_cash=1000000, market=QA.MARKET_TYPE.STOCK_CN,
                 min_trade_money=5000, ratio=0.8,
                 n_up=55, n_down=20, n_short_ma=10, n_long_ma=60, n_atr=20):
        super().__init__(username=username, password=password, code=code, start=start, end=end,
                         frequence=frequence, strategy_id=strategy_id, init_cash=init_cash)
        self.market_type = market
        # ratio交易最大资金比率
        self.ratio = ratio
        # 分别为唐奇安开仓通道.唐奇安平仓通道.短ma.长ma.ATR的参数
        # context.parameter = [55, 20, 10, 60, 20]
        self.N_UP = n_up
        self.N_DN = n_down
        self.N_S_MA = n_short_ma
        self.N_L_MA = n_long_ma
        self.N_ATR = n_atr

        self.end = datetime.datetime.strptime(end, '%Y-%m-%d')
        self.trade_start_time = None
        self.show_Trade_log = show_trade
        self.min_trade_money = min_trade_money
        self.current_percent = 0
        pre_end = datetime.datetime.strptime(start, '%Y-%m-%d')
        if frequence == QA.FREQUENCE.DAY:
            # pre_start = pre_end + datetime.timedelta(days=-N*2)
            pre_start = QA.QA_util_get_pre_trade_date(pre_end, n=n_long_ma + 1)
        else:
            if frequence is QA.FREQUENCE.SIXTY_MIN:
                count = 4
            elif frequence is QA.FREQUENCE.THIRTY_MIN:
                count = 8
            elif frequence is QA.FREQUENCE.FIFTEEN_MIN:
                count = 16
            elif frequence is QA.FREQUENCE.FIVE_MIN:
                count = 48
            else:
                count = 240
            pre_n = int(n_long_ma / count) + 1
            pre_start = QA.QA_util_get_pre_trade_date(pre_end, n=pre_n)
        self.start = datetime.datetime.strptime(pre_start, '%Y-%m-%d')
        self.pre_data = QA.QA_quotation(code, pre_start, pre_end.strftime('%Y-%m-%d'),
                                        frequence=frequence, market=self.market_type,
                                        source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATASTRUCT)
        self.all_data = QA.QA_quotation(code, pre_start, end,
                                        frequence=frequence, market=self.market_type,
                                        source=QA.DATASOURCE.MONGO, output=QA.OUTPUT_FORMAT.DATASTRUCT)

    # sim account on bar
    def on_bar(self, bars):
        bar = bars[0]


    def on_bar_backtest(self, bar):
        # risk需要时间和代码
        bar['datetime'] = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
        bar['code'] = bar.name[1]
        # 计算唐奇安开仓和平仓通道, ATR
        ind = self.market_data.groupby(level=1, sort=False).apply(QA.QA_indicator_Turtle,
                                                                  N_UP=self.N_UP,
                                                                  N_DN=self.N_DN,
                                                                  N_S_MA=self.N_S_MA,
                                                                  N_L_MA=self.N_L_MA,
                                                                  N_ATR=self.N_ATR)
        if self.market_data.shape[0] < 2:
            return
        ind_last_bar = ind.iloc[-2]
        ind_current_bar = ind.iloc[-1]
        upper_band = ind_last_bar.UP
        lower_band = ind_last_bar.DN
        atr = ind_last_bar.ATR
        # 计算开仓的资金比例
        percent = self.ratio
        pos = self.acc.get_position(self.code)
        # hold_vol = self.acc.hold.get(self.code, 0)
        hold_vol = pos.volume_long

        # sell all at end time
        if bar.name[0].day == self.end.day and bar.name[0].month == self.end.month and bar.name[0].year == self.end.year:
            if hold_vol > 0:
                if self.frequence is QA.FREQUENCE.DAY:
                    self.sell(bar, hold_vol, percent=0.0)
                else:
                    if bar.name[0].hour == 15:  # min bar sell at 15:00
                        self.sell(bar, hold_vol, percent=0.0)
        else:
            # 若没有仓位则开仓
            if hold_vol == 0:
                dif = ind_current_bar.DIFF
                # 获取当前价格
                # 上穿唐奇安通道且短ma在长ma上方则开多仓
                if bar.close > upper_band and (dif > 0):
                    self.order_target_percent(bar, position=pos, percent=percent)
                    # print(self.code, '市价单开多仓到比例: ', percent)
            else:
                # 价格跌破唐奇安平仓通道全平仓位止损
                if bar.close < lower_band:
                    self.sell(bar, amount=hold_vol, percent=0.0)
                    # print(self.code, '市价单全平仓位')
                else:
                    # 获取持仓均价
                    vwap = pos.position_price_long
                    # 获取持仓的资金
                    money = pos.open_cost_long
                    # 获取平仓的区间
                    band = vwap - np.array([200, 2, 1.5, 1, 0.5, -100]) * atr
                    grid_percent = float(pd.cut([bar.close], band, labels=[0, 0.25, 0.5, 0.75, 1])[0]) * percent
                    # 选择现有百分比和区间百分比中较小的值(避免开仓)
                    target_percent = np.minimum(money / self.acc.init_cash, grid_percent)
                    # target_percent = grid_percent  # 允许开仓

                    if target_percent != self.current_percent:
                        self.order_target_percent(bar, position=pos, percent=target_percent)

    def order_target_percent(self, bar, position, percent=0.8):
        target_money = self.init_cash * percent
        res_money = target_money - position.open_cost_long
        if res_money > 0:
            buy_money = np.minimum(res_money, self.acc.cash_available)
            buy_amount = int(buy_money / bar.close / 100.25) * 100
            if self.enable_trade(buy_amount * bar.close):
                if self.trade_start_time is None:
                    self.trade_start_time = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
                self.buy(bar, buy_amount, percent)
                self.current_percent = percent
        elif res_money < 0:
            sell_money = -res_money
            sell_amount = int(sell_money / bar.close / 100.0) * 100
            if self.enable_trade(sell_amount * bar.close):
                if self.trade_start_time is None:
                    self.trade_start_time = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
                if sell_amount <= self.acc.sell_available.get(self.code, 0):  # min bar 当日没settle
                    self.sell(bar, sell_amount, percent)
                    self.current_percent = percent

    def enable_trade(self, money):
        if abs(money) >= self.min_trade_money:
            return True
        else:
            return False

    def buy(self, bar, amount, percent):
        order = self.acc.send_order(
            code=self.code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.BUY,
            price=bar.close,
            order_model=QA.ORDER_MODEL.CLOSE,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            print(bar.name[0], 'BUY TO PER', percent, bar.close, amount,
              int(self.acc.cash_available), int(self.acc.cash_available + self.acc.hold.get(self.code, 0) * bar.close))

    def sell(self, bar, amount, percent):
        order = self.acc.send_order(
            code=self.code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.SELL,
            price=bar.close,
            order_model=QA.ORDER_MODEL.MARKET,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            print(bar.name[0], 'SEL TO PER', percent, bar.close, amount, int(self.acc.cash_available),
                  int(self.acc.cash_available + self.acc.hold.get(self.code, 0) * bar.close))

    def plot_indicators(self):
        data = self.all_data.groupby(level=1, sort=False).apply(QA.QA_indicator_Turtle,
                                                                  N_UP=self.N_UP,
                                                                  N_DN=self.N_DN,
                                                                  N_S_MA=self.N_S_MA,
                                                                  N_L_MA=self.N_L_MA,
                                                                  N_ATR=self.N_ATR)
        data = data.xs(self.code, level=1)
        register_matplotlib_converters()
        figure = plt.figure()
        plt.title('{}---'.format(self.code))
        axes1 = figure.add_subplot(411)
        axes2 = figure.add_subplot(412)
        axes3 = figure.add_subplot(413)
        axes4 = figure.add_subplot(414)
        axes1.plot(data.index, data['CLOSE'])
        axes1.plot(data.index, data['UP'])
        axes1.plot(data.index, data['DN'])
        axes2.plot(data.index, data['S_MA'])
        axes2.plot(data.index, data['L_MA'])
        axes3.plot(data.index, data['ATR'])
        axes4.plot(data.index, data['VWAP'])
        axes4.plot(data.index, data['UP'])
        axes4.plot(data.index, data['DN'])

        # left, right = axes1.get_xlim()
        # axes1.hlines(y=box.high, xmin=left, xmax=right, linestyles='dashed')
        multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

        plt.show()